#include <stdio.h>
#include <stdlib.h>

#include <assert.h>
#include <math.h>

#include <mpi.h>

#include "utils.h"
#include "mv.h"

//#define DEBUG

void mv(base_t **A, int nrows, int ncols, int nrows_a_loc, int ncols_a_loc,
        base_t *x, int nrows_x_loc,
        base_t *b, int ncols_b_loc)
{
    int comm_rank, comm_size, ret;

    // query communicator parameters
    MPI_Comm_rank(MPI_COMM_WORLD, &comm_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &comm_size);

#ifdef DEBUG
    printf("(ncols, nrows_x_loc, nrows, ncols_b_loc) for p_%d: (%d, %d, %d, %d)\n", comm_rank, ncols, nrows_x_loc, nrows, ncols_b_loc);
    for (int i = 0; i < nrows_x_loc; i++) {
        printf("xloc for p_%d: %.2f @ index %d\n", comm_rank, x[i], i);
    }
#endif

    // formula for dimensions of matrix mult:
    // h x i * i x j = h x j
    // we need to use the amount of matrix columns (total amount of
    // x vector rows!)
    base_t x_full[ncols];

    // actually only local b is to be computed
    // we need to use the amount of matrix rows (total amount of
    // b vector rows!)
    //base_t b_full[nrows];

    // assume there is no equal split among elements, like n/p for all p_i
    // except for the last process p_{comm_rank - 1}, such that counts are
    // not known without prior communication
    int send_counts[comm_size];
    ret = MPI_Allgather(
        // send buffer, count and type
        &nrows_x_loc, 1, MPI_INT,
        // recv buffer, count and type
        send_counts, 1, MPI_INT,
        // communicator
        MPI_COMM_WORLD
    );

#ifdef DEBUG
    // process sends to itself?
    //send_counts[comm_rank] = 0;
    printf("ret for allgather for p_%d: %d\n", comm_rank, ret);
#endif

    int recv_displacements_loc;
    int recv_displacements[comm_size];
    // perform exclusive prefix sum to get absolute displacements
    ret = MPI_Exscan(
        // send buffer
        &nrows_x_loc,
        // receive buffer
        &recv_displacements_loc,
        // count and type
        1, MPI_INT,
        // operation
        MPI_SUM,
        // communicator
        MPI_COMM_WORLD
    );

#ifdef DEBUG
    printf("ret for exscan for p_%d: %d\n", comm_rank, ret);
#endif

    // 0th exscan value is undefined, but should be 0!
    if (comm_rank == 0)
        recv_displacements_loc = 0;

    // because exscan does not sync all sums to all processes, we need to
    // perform an additional communication round
    ret = MPI_Allgather(
        // send buffer, count and type
        &recv_displacements_loc, 1, MPI_INT,
        // recv buffer, count and type
        recv_displacements, 1, MPI_INT,
        // communicator
        MPI_COMM_WORLD
    );

#ifdef DEBUG
    printf("ret for allgather for p_%d: %d\n", comm_rank, ret);

    for (int i = 0; i < comm_size; i++) {
        printf("(sendcounts, displacements) for p_%d: (%d, %d) @ index %d\n", comm_rank, send_counts[i], recv_displacements[i], i);
    }
#endif

    // make call to allgatherv to gather x vectors for every process
    ret = MPI_Allgatherv(
        // send buffer, count and type
        x, nrows_x_loc, MPI_DOUBLE,
        // recv buffer, count(s), displacements and type
        x_full, send_counts, recv_displacements, MPI_DOUBLE,
        // communicator
        MPI_COMM_WORLD
    );

#ifdef DEBUG
    printf("ret for allgatherv for p_%d: %d\n", comm_rank, ret);

    for (int i = 0; i < ncols; i++) {
        printf("consolidated for p_%d: %.2f @ index %d\n", comm_rank, x_full[i], i);
    }
#endif

    // perform actual matrix vector computation
    for (int i = 0; i < ncols_b_loc; i++) {
        for (int j = 0; j < ncols; j++) {
            b[i] += A[i][j] * x_full[j];
        }
    }
}
